#!/usr/bin/env python3
"""
本地数据库迁移脚本：添加 last_started_at 字段
用于记录任务最后一次开始执行的时间
"""

import pymysql
import os
import traceback
from datetime import datetime
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 数据库连接配置 - 从环境变量读取
def get_db_config():
    """从环境变量获取数据库配置"""
    required_vars = ['MYSQL_HOST', 'MYSQL_PORT', 'MYSQL_USER', 'MYSQL_PASSWORD', 'MYSQL_DATABASE']
    missing_vars = [var for var in required_vars if not os.getenv(var)]
    if missing_vars:
        raise ValueError(f"缺少必需的数据库环境变量: {', '.join(missing_vars)}")
    
    return {
        'host': os.getenv('MYSQL_HOST'),
        'port': int(os.getenv('MYSQL_PORT')),
        'user': os.getenv('MYSQL_USER'),
        'password': os.getenv('MYSQL_PASSWORD'),
        'database': os.getenv('MYSQL_DATABASE'),
        'charset': 'utf8mb4'
    }

def log_message(level, message):
    """简单的日志函数"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {level}: {message}")

def migrate_last_started_at():
    """执行 last_started_at 字段迁移"""
    connection = None
    
    try:
        log_message("INFO", "开始连接数据库")
        db_config = get_db_config()
        connection = pymysql.connect(**db_config)
        cursor = connection.cursor()
        
        log_message("INFO", "开始执行 last_started_at 字段迁移")
        
        # 检查字段是否已存在
        try:
            cursor.execute("SELECT last_started_at FROM tasks LIMIT 1")
            log_message("INFO", "last_started_at 字段已存在，跳过添加字段步骤")
        except pymysql.err.OperationalError as e:
            if "Unknown column" in str(e):
                log_message("INFO", "添加 last_started_at 字段到 tasks 表")
                cursor.execute("ALTER TABLE tasks ADD COLUMN last_started_at DATETIME NULL COMMENT '最后开始执行时间'")
                connection.commit()
                log_message("INFO", "成功添加 last_started_at 字段")
            else:
                raise e
        
        # 为现有任务设置初始值
        log_message("INFO", "开始为现有任务设置 last_started_at 初始值")
        
        # 获取所有任务ID
        cursor.execute("SELECT id, name FROM tasks")
        tasks = cursor.fetchall()
        updated_count = 0
        
        for task_id, task_name in tasks:
            # 获取该任务最近的执行记录
            cursor.execute("""
                SELECT started_at 
                FROM task_executions 
                WHERE task_id = %s AND started_at IS NOT NULL
                ORDER BY started_at DESC 
                LIMIT 1
            """, (task_id,))
            
            result = cursor.fetchone()
            if result and result[0]:
                latest_started_at = result[0]
                cursor.execute("""
                    UPDATE tasks 
                    SET last_started_at = %s 
                    WHERE id = %s
                """, (latest_started_at, task_id))
                updated_count += 1
                log_message("DEBUG", f"任务 {task_id} ({task_name}) 设置 last_started_at = {latest_started_at}")
        
        connection.commit()
        log_message("INFO", f"成功为 {updated_count} 个任务设置了 last_started_at 初始值")
        
        # 验证结果
        log_message("INFO", "验证迁移结果:")
        cursor.execute("SELECT COUNT(*) FROM tasks WHERE last_started_at IS NOT NULL")
        tasks_with_last_started = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM tasks")
        total_tasks = cursor.fetchone()[0]
        
        cursor.execute("""
            SELECT COUNT(DISTINCT task_id) 
            FROM task_executions 
            WHERE started_at IS NOT NULL
        """)
        tasks_with_executions = cursor.fetchone()[0]
        
        log_message("INFO", f"总任务数: {total_tasks}")
        log_message("INFO", f"有执行记录的任务数: {tasks_with_executions}")
        log_message("INFO", f"已设置 last_started_at 的任务数: {tasks_with_last_started}")
        
        if tasks_with_last_started == tasks_with_executions:
            log_message("INFO", "✅ 迁移成功！所有有执行记录的任务都已设置 last_started_at")
        else:
            log_message("WARNING", f"⚠️ 部分任务未设置 last_started_at，可能需要检查")
        
        return True
        
    except Exception as e:
        log_message("ERROR", f"迁移失败: {str(e)}")
        log_message("ERROR", f"错误详情: {traceback.format_exc()}")
        if connection:
            connection.rollback()
        return False
    finally:
        if connection:
            connection.close()

def main():
    """主函数"""
    log_message("INFO", "=" * 50)
    log_message("INFO", "开始执行 last_started_at 字段迁移")
    log_message("INFO", "=" * 50)
    
    success = migrate_last_started_at()
    
    if success:
        log_message("INFO", "=" * 50)
        log_message("INFO", "✅ 迁移完成！")
        log_message("INFO", "=" * 50)
        exit(0)
    else:
        log_message("ERROR", "=" * 50)
        log_message("ERROR", "❌ 迁移失败！")
        log_message("ERROR", "=" * 50)
        exit(1)

if __name__ == "__main__":
    main()